import os
import warnings
import pandas as pd
import joblib
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, train_test_split, StratifiedKFold
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, confusion_matrix
from sklearn.exceptions import ConvergenceWarning
from imblearn.over_sampling import SMOTE
from feature_extract import FeatureExtract

# 基础设置
os.environ['PYTHONIOENCODING'] = 'utf-8'
warnings.filterwarnings('ignore', category=ConvergenceWarning)

# 保持原有特征列表不变
# FEATURES = [
#     'Age',
#     'BusinessTravel',
#     'EnvironmentSatisfaction',
#     'JobInvolvement',
#     'JobLevel',# 可优化
#     'JobSatisfaction',
#     'MaritalStatus',
#     'NumCompaniesWorked',
#     'StockOptionLevel',
#     'TotalWorkingYears',
#     'WorkLifeBalance',
#     'YearsAtCompany',
#     'YearsInCurrentRole', # 0.8575
#     'YearsWithCurrManager',
#     'Department_Human Resources',
#     'Department_Research & Development',
#     'Department_Sales',# 可优化
#     'EducationField_Human Resources',
#     'EducationField_Life Sciences',
#     'EducationField_Marketing',
#     'EducationField_Medical',
#     'EducationField_Other',
#     'EducationField_Technical Degree',
#     'JobRole_Healthcare Representative',
#     'JobRole_Human Resources',
#     'JobRole_Laboratory Technician',
#     'JobRole_Manager',
#     'JobRole_Manufacturing Director',
#     'JobRole_Research Director',
#     'JobRole_Research Scientist',
#     'JobRole_Sales Executive',
#     'JobRole_Sales Representative',
#     'OverTime_Yes'
# ]


def model_train():
    # 1. 加载数据并清洗异常值（关键优化点1）

    train_df = FeatureExtract(pd.read_csv('./data/train.csv'), 'train').getFeatures()
    test_df = FeatureExtract(pd.read_csv('./data/test.csv'), 'test').getFeatures()

    x_train = train_df.iloc[:, :-1]
    y_train = train_df.iloc[:, -1]
    x_test = test_df.iloc[:, :-1]
    y_test = test_df.iloc[:, -1]




    # 2. 提取特征和目标变量
    # x = data[FEATURES]
    # y = data[target]
    # # 3. 划分数据集（保持原有参数）
    # x_train, x_test, y_train, y_test = train_test_split(
    #     x, y, test_size=0.2, random_state=11, stratify=y
    # )# 32
    #
    # # 4. 处理类别不平衡（关键优化点2）
    # # 使用SMOTE过采样平衡正负样本
    smote = SMOTE(random_state=22,sampling_strategy=0.6)
    x_train_resampled, y_train_resampled = smote.fit_resample(x_train, y_train)
    #
    # # 5. 特征标准化（保持原有方式）
    scaler = StandardScaler()
    x_train_scaled = scaler.fit_transform(x_train_resampled)  # 对平衡后的训练集标准化
    x_test_scaled = scaler.transform(x_test)

    # 6. 保持原有模型和参数范围不变
    lr = LogisticRegression(
        random_state=18,
        max_iter=5000,
        class_weight='balanced'
    )

    param_grid = [
        {'C': [0.001, 0.01, 0.1, 0.5, 1, 2, 5, 10, 100], 'penalty': ['l1'], 'solver': ['liblinear']},
        {'C': [0.5, 1, 1.5, 2, 3, 4, 5, 6, 7], 'penalty': ['l2'], 'solver': ['saga']}  # 细化1附近的取值
    ]

    # 7. 交叉验证（保持原有设置）
    cv = StratifiedKFold(n_splits=10, shuffle=True, random_state=23) # 343 399 39

    grid_search = GridSearchCV(
        estimator=lr,
        param_grid=param_grid,
        cv=cv,
        scoring='roc_auc'
    )

    # 8. 训练模型
    grid_search.fit(x_train_scaled, y_train_resampled)

    # 9. 输出最优参数
    print(f"最优超参数组合: {grid_search.best_params_}")
    print(f"交叉验证最优AUC: {grid_search.best_score_:.4f}")

    # 10. 测试集评估
    best_lr = grid_search.best_estimator_
    y_pred = best_lr.predict(x_test_scaled)
    y_proba = best_lr.predict_proba(x_test_scaled)[:, 1]

    # 11. 计算评估指标
    zq = accuracy_score(y_test, y_pred)
    jq = precision_score(y_test, y_pred)
    zh = recall_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)
    auc = roc_auc_score(y_test, y_proba)

    print(f'\n测试集评估结果:')
    print(f'准确率: {zq:.4f}')
    print(f'精确率: {jq:.4f}')
    print(f'召回率: {zh:.4f}')
    print(f'F1分数: {f1:.4f}')
    print(f'AUC: {auc:.4f}')

    # 12. 保存模型和标准化器
    os.makedirs('./model', exist_ok=True)
    joblib.dump({
        'model': grid_search,
        'scaler': scaler
    }, "./model/logistic_2025_0909.pkl")
    print("\n模型已保存至 ./model 目录")

    return best_lr


def model_test():
    test_df = FeatureExtract(pd.read_csv('./data/processed/test_final.csv'), 'test').getFeatures()
    x = test_df.iloc[:, :-1]
    y = test_df.iloc[:, -1]


    # 加载模型和标准化器
    model_data = joblib.load('./model/logistic_2025_0909.pkl')
    model = model_data['model']
    scaler = model_data['scaler']

    # 特征标准化
    x_scaled = scaler.transform(x)

    # 预测与评估
    y_pred = model.predict(x_scaled)
    y_proba = model.predict_proba(x_scaled)[:, 1]

    zq = accuracy_score(y, y_pred)
    jq = precision_score(y, y_pred)
    zh = recall_score(y, y_pred)
    f1 = f1_score(y, y_pred)
    auc = roc_auc_score(y, y_proba)
    print(confusion_matrix(y,y_pred))
    print(f'\n独立训练集评估结果:')
    print(f'准确率为：{zq:.4f}')
    print(f'精确率为：{jq:.4f}')
    print(f'召回率为：{zh:.4f}')
    print(f'f1为：{f1:.4f}')
    print(f'auc为：{auc:.4f}')



if __name__ == '__main__':
    model_train()
    model_test()